#!/usr/bin/env python3
"""
扩展LIBERO数据集，添加多视角图像和相机参数
"""

import os
import h5py
import numpy as np
import argparse
import random
from tqdm import tqdm
import mujoco
import xml.etree.ElementTree as ET
from scipy.spatial.transform import Rotation as R

def get_camera_params(model, data, camera_name):
    """
    从MuJoCo模型中提取相机参数
    返回viewmat和K矩阵
    """
    # 获取相机ID
    camera_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, camera_name)
    if camera_id == -1:
        raise ValueError(f"Camera {camera_name} not found in model")
    
    # 获取相机位置和方向
    cam_pos = np.zeros(3)
    cam_quat = np.zeros(4)
    mujoco.mjf_getCamera(model, data, camera_id, cam_pos, cam_quat)
    
    # 计算视图矩阵 (viewmat)
    rot = R.from_quat([cam_quat[1], cam_quat[2], cam_quat[3], cam_quat[0]])  # MuJoCo使用w,x,y,z顺序
    rot_matrix = rot.as_matrix()
    
    # 构建视图矩阵 (4x4)
    viewmat = np.eye(4)
    viewmat[:3, :3] = rot_matrix
    viewmat[:3, 3] = cam_pos
    
    # 计算相机内参矩阵 (K)
    fovy = model.cam_fovy[camera_id]
    width = model.vis.global_.width
    height = model.vis.global_.height
    
    # 计算焦距
    f = height / (2 * np.tan(fovy * np.pi / 360))
    
    # 构建内参矩阵
    K = np.array([
        [f, 0, width/2],
        [0, f, height/2],
        [0, 0, 1]
    ])
    
    return viewmat, K

def create_random_camera(model, base_camera_name, max_angle_deviation=30):
    """
    基于现有相机创建随机视角的相机
    """
    # 获取基础相机ID
    base_camera_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, base_camera_name)
    if base_camera_id == -1:
        raise ValueError(f"Base camera {base_camera_name} not found in model")
    
    # 获取基础相机位置和方向
    base_pos = np.array(model.cam_pos[base_camera_id*3:(base_camera_id+1)*3])
    base_quat = np.array(model.cam_quat[base_camera_id*4:(base_camera_id+1)*4])
    
    # 创建随机旋转
    angle_x = random.uniform(-max_angle_deviation, max_angle_deviation) * np.pi / 180
    angle_y = random.uniform(-max_angle_deviation, max_angle_deviation) * np.pi / 180
    angle_z = random.uniform(-max_angle_deviation, max_angle_deviation) * np.pi / 180
    
    # 创建随机旋转矩阵
    rot_x = R.from_euler('x', angle_x)
    rot_y = R.from_euler('y', angle_y)
    rot_z = R.from_euler('z', angle_z)
    
    # 组合旋转
    base_rot = R.from_quat([base_quat[1], base_quat[2], base_quat[3], base_quat[0]])  # MuJoCo使用w,x,y,z顺序
    new_rot = base_rot * rot_x * rot_y * rot_z
    
    # 随机位置偏移 (小范围)
    pos_offset = np.random.uniform(-0.1, 0.1, 3)
    new_pos = base_pos + pos_offset
    
    # 转换为四元数
    new_quat = new_rot.as_quat()  # x,y,z,w顺序
    new_quat_mj = np.array([new_quat[3], new_quat[0], new_quat[1], new_quat[2]])  # 转换为MuJoCo的w,x,y,z顺序
    
    return new_pos, new_quat_mj

def add_random_camera_to_xml(xml_str, camera_name, pos, quat):
    """
    向XML添加随机相机
    """
    root = ET.fromstring(xml_str)
    
    # 查找worldbody
    worldbody = root.find(".//worldbody")
    if worldbody is None:
        raise ValueError("worldbody not found in XML")
    
    # 创建新相机元素
    camera = ET.SubElement(worldbody, "camera")
    camera.set("name", camera_name)
    camera.set("pos", f"{pos[0]} {pos[1]} {pos[2]}")
    camera.set("quat", f"{quat[0]} {quat[1]} {quat[2]} {quat[3]}")
    camera.set("mode", "fixed")
    
    # 返回修改后的XML
    return ET.tostring(root, encoding="utf8").decode("utf8")

def extend_dataset_with_camera_params(dataset_path, output_path, num_random_views=3):
    """
    扩展数据集，添加多视角图像和相机参数
    """
    # 打开原始数据集
    with h5py.File(dataset_path, 'r') as src_file:
        # 创建输出数据集
        with h5py.File(output_path, 'w') as dst_file:
            # 复制数据集属性
            for key, value in src_file.attrs.items():
                dst_file.attrs[key] = value
            
            # 遍历所有演示
            for demo_key in tqdm(src_file.keys()):
                if not demo_key.startswith('demo_'):
                    # 复制非演示数据
                    src_file.copy(demo_key, dst_file)
                    continue
                
                # 获取演示数据
                demo_group = src_file[demo_key]
                
                # 创建新的演示组
                dst_demo_group = dst_file.create_group(demo_key)
                
                # 复制演示属性
                for key, value in demo_group.attrs.items():
                    dst_demo_group.attrs[key] = value
                
                # 复制演示数据
                for key in demo_group.keys():
                    if key != 'obs':
                        demo_group.copy(key, dst_demo_group)
                
                # 获取观测数据
                obs_group = demo_group['obs']
                
                # 创建新的观测组
                dst_obs_group = dst_demo_group.create_group('obs')
                
                # 复制观测属性
                for key, value in obs_group.attrs.items():
                    dst_obs_group.attrs[key] = value
                
                # 复制观测数据
                for key in obs_group.keys():
                    obs_group.copy(key, dst_obs_group)
                
                # 获取原始图像数据
                agentview_images = obs_group['agentview_rgb'][:]
                
                # 创建MuJoCo模型
                # 注意：这里需要获取任务的XML文件，这里假设使用通用的XML
                xml_path = os.path.join(os.path.dirname(dataset_path), 'scene.xml')
                if os.path.exists(xml_path):
                    with open(xml_path, 'r') as f:
                        xml_str = f.read()
                else:
                    # 如果没有XML文件，使用默认的LIBERO XML
                    xml_str = """
                    <mujoco model="libero_scene">
                        <worldbody>
                            <camera mode="fixed" name="agentview" pos="0.5 0 1.35" quat="0.653 0.271 0.271 0.653"/>
                        </worldbody>
                    </mujoco>
                    """
                
                # 为每个时间步添加随机视角
                viewmats = []
                Ks = []
                random_view_images = []
                
                # 创建随机相机
                random_cameras = []
                for i in range(num_random_views):
                    # 创建随机相机参数
                    random_pos, random_quat = create_random_camera(None, "agentview")
                    random_cameras.append((f"random_view_{i}", random_pos, random_quat))
                    
                    # 添加随机相机到XML
                    xml_str = add_random_camera_to_xml(xml_str, f"random_view_{i}", random_pos, random_quat)
                
                # 加载模型
                model = mujoco.MjModel.from_xml_string(xml_str)
                data = mujoco.MjData(model)
                
                # 初始化渲染器
                renderer = mujoco.Renderer(model)
                
                # 为每个时间步生成随机视角图像和相机参数
                for t in range(len(agentview_images)):
                    # 这里需要设置模型状态，但我们没有完整的状态信息
                    # 在实际应用中，您需要从数据集中加载完整的状态
                    # 这里我们只是演示概念
                    
                    # 更新场景
                    mujoco.mj_forward(model, data)
                    
                    # 收集每个随机相机的参数和图像
                    frame_viewmats = []
                    frame_Ks = []
                    frame_images = []
                    
                    for camera_name, _, _ in random_cameras:
                        # 获取相机参数
                        viewmat, K = get_camera_params(model, data, camera_name)
                        frame_viewmats.append(viewmat)
                        frame_Ks.append(K)
                        
                        # 渲染图像
                        renderer.update_scene(data, camera=camera_name)
                        img = renderer.render()
                        frame_images.append(img)
                    
                    viewmats.append(frame_viewmats)
                    Ks.append(frame_Ks)
                    random_view_images.append(frame_images)
                
                # 将数据转换为numpy数组
                viewmats = np.array(viewmats)
                Ks = np.array(Ks)
                random_view_images = np.array(random_view_images)
                
                # 添加到数据集
                dst_obs_group.create_dataset('viewmats', data=viewmats)
                dst_obs_group.create_dataset('Ks', data=Ks)
                
                # 添加随机视角图像
                for i in range(num_random_views):
                    dst_obs_group.create_dataset(f'random_view_{i}_rgb', data=random_view_images[:, i])

def main():
    parser = argparse.ArgumentParser(description='扩展LIBERO数据集，添加多视角图像和相机参数')
    parser.add_argument('--dataset', type=str, required=True, help='原始数据集路径')
    parser.add_argument('--output', type=str, required=True, help='输出数据集路径')
    parser.add_argument('--views', type=int, default=3, help='随机视角数量')
    args = parser.parse_args()
    
    extend_dataset_with_camera_params(args.dataset, args.output, args.views)

if __name__ == '__main__':
    main()